#!/usr/bin/env python
#TODO: posibly replace ping with arping packets.
#Raul Jimenez

"""
Usage: pwol.py [options]
Wake on Lan: automatically wake computers on the same broadcast address during 
a user specified time range. Uses ping to scan for mac addresses and sends the 
hosts a 'MagicPacket'(FF*6 MacAddress*20) once ping cannot reach the hosts.
Options:
  -h, --help            show this help message and exit
  -v, --verbose
  -q, --quiet
  -o, --one-shot
  -g, --gui
  -f MAC_FILE, --file=MAC_FILE
  -l LOG, --log=LOG
  -d DELAY, --delay=DELAY
  -p, --ping
  -w PING_WAIT, --ping-wait=PING_WAIT
  -W WAKE_WAIT, --wake-wait=WAKE_WAIT
"""

import subprocess,re,socket,struct,threading,Queue,time,sys,optparse,gtk,gobject
from Queue import Queue

gobject.threads_init()

#parse Options
parser=optparse.OptionParser()

#Default User Settings
parser.set_defaults(debug=True)
parser.set_defaults(one_shot=False)
parser.set_defaults(mac_file="./mac_list")
parser.set_defaults(log="./Wol_Log")
parser.set_defaults(delay=120)
parser.set_defaults(ping=False)
parser.set_defaults(ping_wait=0.1)
parser.set_defaults(wake_wait=0.1)

#Wether to display debug information.
parser.add_option("-v","--verbose",action='store_true',dest="debug")
parser.add_option("-q","--quiet",action='store_false',dest="debug")
#Do a single run (one shot mode).
parser.add_option("-o","--one-shot",action='store_true',dest="one_shot")
#Open in gui mode.
parser.add_option("-g","--gui",action='store_false',dest="one_shot")
#Location to save MacAddresses
parser.add_option("-f","--file",dest="mac_file")
#Destination to log file, omit to not log.
parser.add_option("-l","--log",dest="log")
#Delay in seconds between Watch Dog cycles.
parser.add_option("-d","--delay",type="float",dest="delay")
#Wether to try and discover more hosts.
parser.add_option("-p","--ping",action='store_true',dest="ping")
#Delay in seconds between pings.
parser.add_option("-w","--ping-wait",type="float",dest="ping_wait")
#Delay in seconds between wakes.
parser.add_option("-W","--wake-wait",type="float",dest="wake_wait")
options=parser.parse_args()[0]

#Global Variables:
mac_list=[]              #[mac,host,alive] alive: 0=false, 1=true, 3=error
if not options.one_shot:
    mac_list=gtk.ListStore(str,str,str)
mac_list_lock=None
exit=False
file_length=0


class Interface:
    """Creates the window for gui mode."""

    def __init__(self, list, dog):
        """Create all the widgets."""
        self.list=list
        self.tree=None
        self.watch_dog=dog
        print self.watch_dog, self.watch_dog.end_time
        (self.adj,self.spinners)=self.make_spinners()
        self.current_time=gtk.Label("Current Time: %02i:%02i:%02i"%(1,6,4))
        self.buttons=self.make_buttons()
        self.hbox=self.make_hbox()
        self.vbox=self.make_vbox()
        self.window=self.make_window()

    def ping(self, widget=None, event=None, data=None):
        print "ping!"
        #import pdb
        #pdb.set_trace()
        (model, selected) = self.tree.get_selection().get_selected_rows()
        for selection in selected:
            host = model.get_value(model.get_iter(selection), 1)
            stat=Ping()

    def quit(self, widget=None, event=None, data=None):
        """Signal to exit."""
        global exit
        gtk.main_quit()
        exit=True
        print "Exiting.."

    def set_time(self,adj,time_remaining):
        """Update range of time for watch dog to remain active."""
        print "Time from %i, to %i!"%(time_remaining[0][time_remaining[1]],adj.value)
        time_remaining[0][time_remaining[1]]=adj.value

    def make_tree(self):
        """Create tree display."""
        col=[]
        cell=gtk.CellRendererText()
        col.append(gtk.TreeViewColumn('Mac Address',cell,text=0))
        col.append(gtk.TreeViewColumn('Host',cell,text=1))
        col.append(gtk.TreeViewColumn('Status',cell,text=2))
        if not self.list: self.list=gtk.ListStore(str,str,str)
        self.tree=gtk.TreeView(self.list)
        for x in col:
            self.tree.append_column(x)
            x.set_visible(True)
        self.tree.set_search_column(0)
        col[0].set_sort_column_id(0)
        self.tree.set_reorderable(False)
        scroll=gtk.ScrolledWindow()
        scroll.set_policy(gtk.POLICY_AUTOMATIC,gtk.POLICY_ALWAYS)
        scroll.add_with_viewport(self.tree)
        scroll.show()
        self.tree.get_selection().set_mode(gtk.SELECTION_MULTIPLE)
        self.tree.show()
        return scroll

    def make_spinners(self):
        """Create time range widgets."""
        spinners=[]
        adj=[]
        adj.append(gtk.Adjustment(value=1,lower=1,upper=24,step_incr=1,page_incr=5,page_size=0))
        adj[0].connect("value_changed",self.set_time,[self.watch_dog.start_time,0])
        adj.append(gtk.Adjustment(value=0,lower=0,upper=59,step_incr=1,page_incr=5,page_size=0))
        adj[1].connect("value_changed",self.set_time,[self.watch_dog.start_time,1])
        adj.append(gtk.Adjustment(value=24,lower=1,upper=24,step_incr=1,page_incr=10,page_size=0))
        adj[2].connect("value_changed",self.set_time,[self.watch_dog.end_time,0])
        adj.append(gtk.Adjustment(value=0,lower=0,upper=59,step_incr=1,page_incr=10,page_size=0))
        adj[3].connect("value_changed",self.set_time,[self.watch_dog.end_time,1])
        for x in adj: spinners.append(gtk.SpinButton(x,climb_rate=1,digits=0))
        return (adj,spinners)

    def make_buttons(self):
        """Create buttons and assign call backs."""
        buttons=[]
        for x in ['Ping','Ping All','Wake','Wake All','Quit']:
            buttons.append(gtk.Button(x,None,False))
        buttons[0].connect("clicked", self.ping)
        buttons[4].connect("clicked", self.quit)
        return buttons

    def make_hbox(self):
        """Arrange widgets in window in to rows."""
        box=[]
        labels=[gtk.Label("Start [hh,mm]"),gtk.Label("Stop [hh,mm]")]
        hbox=gtk.HBox(False,0)
        hbox.pack_start(self.buttons[0],True,False,0)
        hbox.pack_start(self.buttons[1],True,False,0)
        hbox.pack_start(labels[0],True,False,0)
        hbox.pack_start(self.spinners[0],False,False,0)
        hbox.pack_start(self.spinners[1],False,False,0)
        box.append(hbox)
        hbox=gtk.HBox(False,0)
        hbox.pack_start(self.buttons[2],True,False,0)
        hbox.pack_start(self.buttons[3],True,False,0)
        hbox.pack_start(labels[1],True,False,0)
        hbox.pack_start(self.spinners[2],False,False,0)
        hbox.pack_start(self.spinners[3],False,False,0)
        box.append(hbox)
        hbox=gtk.HBox(False,0)
        hbox.pack_start(self.current_time,True,False,0)
        hbox.pack_start(self.buttons[4],True,False,0)
        box.append(hbox)
        return box

    def make_vbox(self):
        """Arrange widget rows in Window."""
        vbox=gtk.VBox(False,4)
        vbox.pack_start(self.hbox[0],False,False,0)
        vbox.pack_start(self.hbox[1],False,False,0)
        vbox.pack_start(self.hbox[2],False,False,0)
        vbox.pack_start(self.make_tree(),True,True,0)
        return vbox

    def make_window(self):
        """create the window and assign close callback."""
        window=gtk.Window()
        window.connect("delete_event",self.quit)
        window.set_title('Wake On Lan')
        window.set_resizable(True)
        window.resize(350,350)
        window.add(self.vbox)
        window.show_all()
        return window


class Ping(threading.Thread):
    """Separate ping thread. Ping remote host to confirm status and survey for new\nhost computers."""

    def __init__(self,release,args,resource):
        """Pings host and writes status to mac_list[idx][2] if idx is given."""
        threading.Thread.__init__(self,name="Ping Thread #"+str(args[0]))
        self.done=release
        self.host=args[0]
        self.indx=args[1]
        self.start()

    def ping_host(self,bcast=False):
        """Call ping through os."""
        cmd=['ping','-q',"-w2",'-c1']
        stat=3
        if bcast: cmd.append('-b')
        else: cmd.append(self.host)
        log("Pinging %s"%cmd)
        try: stat=int(re.split('[\n\s]*',subprocess.Popen(cmd,stdout=subprocess.PIPE).communicate()[0])[15])
        except Exception,e: log(err="Exception in ping_host %s"%str(e))
        return stat

    def run(self):
        """Call ping and check if mac_list needs to be updated."""
        global exit
        try:
            stat=self.ping_host()
            if self.indx:
                mac_list=lock_list()
                if not mac_list[self.indx][2]==stat:
                    mac_list[self.indx][2]=stat
                release_list()
            self.done()
            log("Ping Status %s, %i"%(self.host,stat))
        except Exception,e: 
            log(err="Exception in Ping %s"%str(e))
            exit=True


class Dispatcher(threading.Thread):
    """Thread to handle Queue containing jobs to dispatch. Takes a function\nfunc(release,queue_element,resource)"""

    def __init__(self,func,amount=None,resource=None, wait=None, name="", one_shot=False):
        """Creates a Dispatcher that calls func on queued jobs, allowing up to 'amount' at\na time. usieng resourece if given and sleeping between calls for\namount 'wait'. Calls func(release,queue_element,resource) Func needs to calls\nrelease() from within."""
        threading.Thread.__init__(self)
        log("thread for %s started. with resource %s"%(name,str(type(resource))))
        self.job_func=func
        self.one_shot=one_shot
        self.resource=resource
        self.queue=Queue()
        self.semaphore=threading.Semaphore(amount)
        self.wait=wait
        self.setName(str(name))
        self.start()

    def queue_job(self,job):
        self.queue.put(job)

    def release(self):
        """Free resources when function finishes executing."""
        self.semaphore.release()

    def run(self):
        """While exit has not been set dispatch jobs."""
        global exit
        try:
            started=False
            while not exit:
                if not self.queue.empty():
                    started=True
                    self.semaphore.acquire()
                    log("Exit %s, Size %i, Name %s"%(exit,self.queue.qsize(),self.getName()))
                    self.job_func(self.release,self.queue.get(),self.resource)
                    time.sleep(self.wait)
                else: 
                    if self.one_shot and started: return
                    
                    time.sleep(.5)
        except Exception,e: 
            log(err="Exception in Dispatch %s"%str(e))
            exit=True


class Watch_Dog(threading.Thread):
    """Monitors remote hosts' status through ping, if a host computer goes off during\nspecified time range, it queues a wake job."""

    def __init__(self,dispatch_ping,dispatch_wake,delay,bcast):
        """Acquires dispatched to ping and wake remote hosts."""
        threading.Thread.__init__(self,name="Watch Dog")
        log("Watch Dog started.")
        (self.d_ping,self.d_wake)=(dispatch_ping,dispatch_wake)
        self.start_time=[24,59]; self.end_time=[24,59]
        self.delay=delay
        self.bcast=bcast
        self.start()

    def wake(self,all=None):
        """Wake sleeping hosts or if 'all' is True wake all hosts."""
        mac_list=lock_list()
        log("Wake_all(%s)"%str(all))
        for host in mac_list:
            if not all or host[2]:
                self.d_wake.queue_job(host[0])
        release_list()

    def run(self):
        """Will ping hosts to check state and parse arp file to gain mac addresses and\nwake hosts not responding to ping."""
        global exit,options

        def within_range(start,end):
            """Check current time to see if it is within start and end."""
            current_time=time.localtime()[3:5]
            if current_time[0] <= end[0] and current_time[0] >= start[0]:
                if (current_time[0]==start[0] and current_time[1] < start[1]) or (current_time[0]==end[0] and current_time[1] > end[1]): return False
                return True
            return False
        try:
            while not exit:
                log("Watch Dog Cycle %i:%i<->%i:%i"%(self.start_time[0],self.start_time[1],self.end_time[0],self.end_time[1]))
                if within_range(self.start_time,self.end_time):
                    print "within"
                    parse_arp()
                    self.wake()
                    ping_all(self.bcast,self.d_ping)
                    time.sleep(self.delay)
                else:
                    time.sleep(self.delay)
        except Exception,e: 
            log(err="Exception in Watch_Dog %s"%str(e))
            exit=True
#Global Functions

def find_macs(item, seq, a=None, b=None, within=0, insert=False):
    """Find the index of a value within, set within to 0=mac,1=host. If insert is set to true it will insert a value not currently in the list(mac only)"""
    log("Length of Sequence %i."%len(seq))
    size=len(seq)
    answer = None
    if size==0:
        if insert and within==0:
            seq.append([item,None,0])
            return 0
        return None
    if a==None: (a,b)=(0,size-1)
    c=a+((b-a)/2)
    log("dif: a%i,b%i,(b-a)%i,c%i"%(a,b,(b-a),c))
    if c==a:
        if item==seq[b][within]:
            return b
        if item==seq[a][within]:
            return a
        else:
            log("Item [%s] not found."%str(item))
            if insert and within==0:
                token=[item,None,0]
                if item<seq[a][within]:
                    seq.insert(a,token)
                    return a
                if item<seq[b][within]:
                    seq.insert(b,token)
                    return b
                else: 
                    r=b+1
                    if r==size: seq.append(token)
                    else: seq.insert(r,token)
                    return r
            return None
    if item > seq[c][within]: 
        return find_macs(item, seq, c, b, within, insert)
    else: 
        return find_macs(item, seq, a, c, within, insert)

def parse_arp():
    """Parse arp list to get MacAddresses and host address."""
    global options
    file=open('/proc/net/arp','r')
    list=[]
    null='00:00:00:00:00:00'
    mac_list=lock_list()
    for line in re.split('\n',file.read()):
        tokens=re.split('[\s]*',line)
        print tokens
        if len(tokens)>5: 
            log("mac[%s],host[%s]"%(tokens[3],tokens[0]))
            mac=''
            in_list=False
            if tokens[3]==null or not len(tokens[3])==len(null): continue
            hex=re.split(':',tokens[3]) 
            mac=hex[0]+hex[1]+hex[2]+hex[3]+hex[4]+hex[5]
            idx=find_macs(mac,mac_list,insert=True)
            log("Found at %i"%idx)
            mac_list[idx][1]=tokens[0]
            mac_list[idx][2]=1
    if options.debug: 
        print "parse arp"
        for x,y,z in mac_list: print "[%s,%s,%s]"%(x,y,z)
    release_list()
    file.close()

def parse_macs(file_name=None):
    """Parse file with list of macs."""
    global options, file_length
    if not file_name: 
        file_name="./mac_list"
        log("Saving MacAddresses to %s"%file_name)
    try: file=open(file_name,'r')
    except: return
    mac_list=lock_list()
    for line in re.split('\n',file.read()):
        if line and len(line)==12: find_macs(line,mac_list,insert=True)
        else: log(err="Invalid MacAdress {%s}"%(line))
    file_length=len(mac_list)
    release_list()
    file.close()
    if options.debug: 
        print "parse macs"
        for x,y,z in mac_list: print "[%s,%s,%s]"%(x,y,z)

def parse_broadcast():
    """Parse broadcast address from /sbin/ifconfig."""
    bcast=''
    for x in re.split('[\n\s:.]*',subprocess.Popen('/sbin/ifconfig',stdout=subprocess.PIPE).communicate()[0])[18:21]: bcast+=x+'.'
    return bcast

def ping_all(broadcast, dispatch_ping):
    """Ping all machines on broadcast address in range 1-254."""
    host_indx=1
    while host_indx<255:
        host=broadcast+str(host_indx)
        found=None
        mac_list=lock_list()
        for x in range(len(mac_list)):
            if mac_list[x][1]==host: found=x
        release_list()
        dispatch_ping.queue_job([host,found])
        host_indx+=1

def wake_host(done,mac,sock,prt=9):
    """Broadcast 'MagicPacket' to wake remote host. (packet=FF*16mac*20)"""
    data=('F'*12)+(mac*20); indx=0; packet=''
    while indx<len(data):
        packet+=struct.pack('B',int(data[indx:indx+2],16)) 
        indx+=2 
    log("Wakeing %s."%mac)
    sock.sendto(packet,('<broadcast>',prt))
    done()

def save_macs(name=None):
    """Save MacAddresses to a file."""
    file_name='./mac_list'
    if name: file_name=name
    file=open(file_name,'w')
    mac_list=lock_list()
    for mac in mac_list: file.write(mac[0]+'\n')
    release_list()
    file.close()
    log("Macs saved to %s"%file_name)

def log(mesg=None, err=None):
    """Log errors/messeges to option.log if loging is on. if debug is on  print \nthem to screen."""
    global options
    file=None
    text=""
    time_str="Time: %s\n"%time.ctime()
    if options.log: file=open(options.log,'a')
    if err:
        text+='\tErr: %s\n'%str(err)
        sys.stderr.write(text)
    if mesg:
        text+='\t%s\n'%str(mesg)
    if options.debug: print text
    if options.log:
        file.write(time_str+text)
        file.close()

def create_socket():
    """Creates socket."""
    sock=socket.socket(socket.AF_INET,socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET,socket.SO_BROADCAST,1)
    return sock

def lock_list():
    """Lock mac_list from being accessed out side of current thread."""
    global mac_list, mac_list_lock
    mac_list_lock.acquire()
    #log("mac_list locked within %s"%threading.currentThread().getName())
    return mac_list

def release_list():
    """Release mac_list."""
    global mac_list, mac_list_lock
    mac_list_lock.release()
    #log("mac_list released from within %s"%threading.currentThread().getName())

def shot(ping=False):
    """Do not loop, wake all once. Ping and parse arp if specified."""
    global mac_list,options
    log("Starting up in one shot mode.")
    parse_macs(options.mac_file)
    if not mac_list:
        ping=True
        log("No mac list found auto discovering hosts.")
    parse_arp()
    ping_disp=None
    wake_disp=None
    if file_length: wake_disp=Dispatcher(func=wake_host,amount=20,wait=0.01,resource=create_socket(), one_shot=True, name="Wake Dispatcher")
    for host in mac_list:
        if wake_disp: wake_disp.queue_job(host[0])
    if ping:
        ping_disp=Dispatcher(func=Ping,amount=5,wait=0.0, one_shot=True, name="Ping Dispatcher")
        ping_all(parse_broadcast(),ping_disp)
        ping_disp.join()
        for t in threading.enumerate():
            try: 
                if not t==threading.currentThread():
                    log("Joining %s"%t.getName())
                    t.join()
                    log("Done Joining %s"%t.getName())
                else: log("Skiped %s"%t.getName())
            except Exception,e: log(err="Exception in shot %s"%str(e))
        parse_arp()
    if wake_disp: wake_disp.join()
#Main Loop
if __name__ == '__main__':
    log("###Session Started.###")
    log(str(options))
    mac_list_lock=threading.RLock()
    if options.one_shot:
        shot(options.ping)
    else:
        parse_macs(options.mac_file)
        parse_arp()
        ping_disp=Dispatcher(func=Ping,amount=5,wait=options.ping_wait)
        wake_disp=Dispatcher(func=wake_host,amount=20,wait=options.wake_wait,resource=create_socket())
        watch_dog=Watch_Dog(ping_disp,wake_disp,options.delay,parse_broadcast())
        gui=Interface(mac_list,watch_dog)
        gtk.main()
        print "Closing threads."
        for thread in threading.enumerate():
            if not thread==threading.currentThread(): 
                log("Joining %s"%thread.getName())
                thread.join()
            else: log("Skiped Thread %s"%thread.getName())
    if options.debug:
        for x,y,z in mac_list: print "[%s,%s,%s]"%(x,y,z)
    save_macs(options.mac_file)
    log("---Session Closed.---\n\n\n")
